import pandas as pd
import numpy as np
from sklearn.ensemble import RandomForestRegressor
import joblib
from datetime import datetime, timedelta
from sklearn.preprocessing import StandardScaler


class SalesPredictor:
    def __init__(self, model_path='sales_model.pkl'):
        self.model_path = model_path
        self.model = None
        self.cutoff_date = None
        self.max_days = 90  # 商品生命周期90天
        self.forecast_days = 14  # 预测未来14天
        self.history_days = 30  # 使用最近30天的历史数据作为特征
        self.scaler = None  # 用于特征标准化

    def prepare_data(self, raw_df, cutoff_date=None):
        """
        准备训练数据：
        1. 过滤超过90天的数据
        2. 补充缺失日期销量为0
        3. 生成时序特征和目标变量
        4. 过滤上市不足44天的商品（30天历史+14天预测）
        """
        # 确保日期格式为YYYY-MM-DD
        raw_df['date_id'] = pd.to_datetime(raw_df['date_id'], format='%Y-%m-%d')
        raw_df['new_date'] = pd.to_datetime(raw_df['new_date'], format='%Y-%m-%d')

        # 应用截止日期过滤
        if cutoff_date:
            self.cutoff_date = pd.to_datetime(cutoff_date, format='%Y-%m-%d')
            raw_df = raw_df[raw_df['date_id'] <= self.cutoff_date]
        else:
            # 如果没有提供截止日期，使用数据中最大日期
            self.cutoff_date = raw_df['date_id'].max()

        # 生成完整日期范围
        full_data = []
        grouped = raw_df.groupby('spu_code')
        mature_spus = []  # 记录上市时间超过44天的SPU

        for spu, group in grouped:
            launch_date = group['new_date'].min()

            # 计算商品在截止日期时的上市天数
            days_since_launch = (self.cutoff_date - launch_date).days + 1

            # 过滤上市不足44天的商品
            if days_since_launch < (self.history_days + self.forecast_days):
                print(
                    f"SPU {spu} skipped: Only {days_since_launch} days since launch (less than {self.history_days + self.forecast_days} days)")
                continue

            mature_spus.append(spu)

            # 计算商品生命周期内的日期范围
            date_range = pd.date_range(
                start=launch_date,
                periods=min(days_since_launch, self.max_days),
                freq='D'
            )

            full_dates = pd.DataFrame({'date_id': date_range, 'spu_code': spu})

            # 合并实际销售数据
            merged = pd.merge(full_dates, group, on=['date_id', 'spu_code'], how='left')
            merged['daily_sale_qty'] = merged['daily_sale_qty'].fillna(0)
            merged['new_date'] = launch_date

            full_data.append(merged)

        if not mature_spus:
            print("Warning: No mature SPUs found (all have less than min days since launch)")
            return np.array([]), np.array([])

        full_df = pd.concat(full_data, ignore_index=True)

        # 生成训练样本
        X, y = [], []
        for spu, group in full_df.groupby('spu_code'):
            sales_series = group.sort_values('date_id')['daily_sale_qty'].values

            # 滑动窗口生成样本
            for i in range(self.history_days, len(sales_series) - self.forecast_days + 1):
                # 提取历史窗口数据
                history = sales_series[i - self.history_days:i]

                # 时序特征工程
                features = []

                # 1. 基础统计特征
                features.append(np.mean(history))  # 历史平均销量
                features.append(np.median(history))  # 历史中位数
                features.append(np.std(history))  # 历史标准差
                features.append(np.min(history))  # 历史最小值
                features.append(np.max(history))  # 历史最大值
                features.append(np.sum(history))  # 历史总销量

                # 2. 时间相关特征
                current_day = i  # 当前天数（从上市开始）
                features.append(current_day)  # 当前天数

                # 3. 移动平均特征
                features.append(np.mean(history[-7:]))  # 最近7天平均
                features.append(np.mean(history[-14:]))  # 最近14天平均
                features.append(np.mean(history[-21:]))  # 最近21天平均

                # 4. 变化趋势特征
                if len(history) > 1:
                    features.append(history[-1] - history[-2])  # 昨日变化
                    features.append(history[-1] - history[-7])  # 周环比变化

                # 5. 指数加权移动平均
                ewma = pd.Series(history).ewm(span=7).mean()
                features.append(ewma.iloc[-1])

                # 6. 季节性特征（星期几）
                # 假设group有日期信息，取最后一天的星期几
                last_date = group['date_id'].iloc[i - 1]
                features.append(last_date.weekday())  # Monday=0, Sunday=6

                # 目标：未来14天累计销量
                future_sales = np.sum(sales_series[i:i + self.forecast_days])

                X.append(features)
                y.append(future_sales)

        print(f"Prepared training data from {len(mature_spus)} mature SPUs, {len(X)} samples")
        return np.array(X), np.array(y)

    def full_train(self, data_path, cutoff_date=None):
        """全量训练模型"""
        # 读取Excel文件
        if data_path.endswith('.xlsx'):
            raw_df = pd.read_excel(data_path, engine='openpyxl')
        elif data_path.endswith('.xls'):
            raw_df = pd.read_excel(data_path)
        else:
            raise ValueError("Unsupported file format. Only .xlsx and .xls files are supported.")

        X, y = self.prepare_data(raw_df, cutoff_date)

        # 检查是否有足够的数据
        if len(X) == 0:
            print("Warning: No training data available after filtering.")
            return None

        # 特征标准化
        self.scaler = StandardScaler()
        X_scaled = self.scaler.fit_transform(X)

        # 训练模型
        self.model = RandomForestRegressor(
            n_estimators=150,
            max_depth=10,
            min_samples_leaf=4,
            random_state=42,
            n_jobs=-1
        )
        self.model.fit(X_scaled, y)

        # 保存模型和标准化器
        joblib.dump({
            'model': self.model,
            'scaler': self.scaler
        }, self.model_path)

        print(f"Model trained on {len(X)} samples and saved to {self.model_path}")
        return self.model

    def incremental_train(self, new_data_path):
        """增量训练模型"""
        # 加载现有模型
        try:
            saved_data = joblib.load(self.model_path)
            self.model = saved_data['model']
            self.scaler = saved_data['scaler']
        except:
            print("No existing model found. Creating new model.")
            self.scaler = StandardScaler()
            self.model = RandomForestRegressor(n_estimators=100, random_state=42)

        # 读取Excel文件
        if new_data_path.endswith('.xlsx'):
            new_df = pd.read_excel(new_data_path, engine='openpyxl')
        elif new_data_path.endswith('.xls'):
            new_df = pd.read_excel(new_data_path)
        else:
            raise ValueError("Unsupported file format. Only .xlsx and .xls files are supported.")

        # 准备增量数据
        X_new, y_new = self.prepare_data(new_df, self.cutoff_date)

        # 检查是否有新数据
        if len(X_new) == 0:
            print("Warning: No new training data available after filtering.")
            return None

        # 增量训练
        # 1. 标准化新数据
        X_new_scaled = self.scaler.transform(X_new)

        # 2. 增量训练模型
        # 创建新的模型集合（原模型+新数据训练的模型）
        new_model = RandomForestRegressor(
            n_estimators=50,
            random_state=42,
            warm_start=False
        )
        new_model.fit(X_new_scaled, y_new)

        # 3. 合并模型
        self.model.estimators_ += new_model.estimators_
        self.model.n_estimators = len(self.model.estimators_)

        # 保存更新后的模型
        joblib.dump({
            'model': self.model,
            'scaler': self.scaler
        }, self.model_path)

        print(f"Model incrementally trained with {len(X_new)} new samples")
        return self.model

    def extract_features(self, sales_series, current_day):
        """从销售序列中提取时序特征"""
        # 确保有足够的历史数据
        if len(sales_series) < self.history_days:
            # 不足则前面补0
            padded_sales = np.pad(sales_series, (self.history_days - len(sales_series), 0), 'constant')
            history = padded_sales
        else:
            history = sales_series[-self.history_days:]

        features = []

        # 1. 基础统计特征
        features.append(np.mean(history))  # 历史平均销量
        features.append(np.median(history))  # 历史中位数
        features.append(np.std(history))  # 历史标准差
        features.append(np.min(history))  # 历史最小值
        features.append(np.max(history))  # 历史最大值
        features.append(np.sum(history))  # 历史总销量

        # 2. 时间相关特征
        features.append(current_day)  # 当前天数

        # 3. 移动平均特征
        features.append(np.mean(history[-7:]))  # 最近7天平均
        features.append(np.mean(history[-14:]))  # 最近14天平均
        features.append(np.mean(history[-21:]))  # 最近21天平均

        # 4. 变化趋势特征
        if len(history) > 1:
            features.append(history[-1] - history[-2])  # 昨日变化
            if len(history) > 7:
                features.append(history[-1] - history[-7])  # 周环比变化
            else:
                features.append(0)
        else:
            features.extend([0, 0])

        # 5. 指数加权移动平均
        ewma = pd.Series(history).ewm(span=7).mean()
        features.append(ewma.iloc[-1])

        # 6. 季节性特征（星期几）
        # 假设最后一天是今天
        last_date = datetime.now()  # 实际应用中应根据数据确定
        features.append(last_date.weekday())  # Monday=0, Sunday=6

        return np.array(features)

    def predict(self, spu_data):
        """预测单个SPU未来14天销量"""
        # 确保模型已加载
        if self.model is None or self.scaler is None:
            try:
                saved_data = joblib.load(self.model_path)
                self.model = saved_data['model']
                self.scaler = saved_data['scaler']
            except:
                raise ValueError("Model not trained or loaded. Please train a model first.")

        # 确保输入是DataFrame
        if not isinstance(spu_data, pd.DataFrame):
            raise ValueError("Input must be a pandas DataFrame")

        # 检查数据是否为空
        if len(spu_data) == 0:
            print("Warning: Input data is empty.")
            return 0

        # 转换日期格式
        spu_data['date_id'] = pd.to_datetime(spu_data['date_id'], format='%Y-%m-%d')
        spu_data['new_date'] = pd.to_datetime(spu_data['new_date'], format='%Y-%m-%d')

        # 获取上市日期和最后观测日期
        launch_date = spu_data['new_date'].min()
        last_observed_date = spu_data['date_id'].max()
        current_day = (last_observed_date - launch_date).days + 1  # 当前是第几天（从1开始）

        # 检查是否超过预测范围
        if current_day > self.max_days - self.forecast_days:
            spu_code = spu_data['spu_code'].iloc[0] if 'spu_code' in spu_data.columns else 'Unknown'
            print(
                f"SPU {spu_code} exceeds {self.max_days}-day lifecycle for {self.forecast_days}-day forecast (current day: {current_day}), no prediction made")
            return 0

        # 按日期排序
        spu_data = spu_data.sort_values('date_id')
        sales_series = spu_data['daily_sale_qty'].values

        # 提取时序特征
        features = self.extract_features(sales_series, current_day)

        # 标准化特征
        features_scaled = self.scaler.transform([features])

        # 预测
        try:
            prediction = self.model.predict(features_scaled)[0]
            return max(0, prediction)  # 确保非负预测
        except Exception as e:
            print(f"Prediction failed: {str(e)}")
            return 0


# 示例用法
if __name__ == "__main__":
    predictor = SalesPredictor()

    # 全量训练
    # predictor.full_train('sales.xlsx')

    # 模拟不同阶段的预测
    sales_array = [1977,146,63,39,38,33,33,44,42,29,33,28,50,23,16,44,19,15,13,20,16,12,13,22,18,19,19,12,22,20]

    for i in range(len(sales_array)):  # 从第7天开始预测
        sub_sales_array = sales_array[:i + 1]
        known_days = len(sub_sales_array)
        new_spu_data = pd.DataFrame({
            'spu_code': ['LO2412090'] * known_days,
            'date_id': pd.date_range(start='2025-01-10', periods=known_days, freq='D'),
            'new_date': ['2025-01-10'] * known_days,
            'daily_sale_qty': sub_sales_array
        })

        # 执行预测
        forecast = predictor.predict(new_spu_data)
        print(f"第{i + 1:02}天销量预测: {forecast:.2f} (基于最近{min(known_days, 30)}天数据)")